import hydra
import glob
import trimesh
import open3d as o3d
import numpy as np
import pickle as pkl
import os.path as osp
import mindspore as ms
import mindspore.Tensor as Tensor


class XHumansDataSet():
    """
    Construct SMPL-X dataset from
    Expressive Body Capture: 3D Hands, Face, and Body from a Single Image
    https://arxiv.org/pdf/1904.05866.pdf

    Args:
        dataset_path: path of the dataset
        model_type: type of the model
        use_pca: use pca or not
        num_pca_comps: number of pca components
        flat hand mean: mean of flat hand

    Return:
        data(Tensor)

    Supported Platforms:
        GPU
    """

    def __init__(self, dataset_path, model_type, use_pca,
                 num_pca_comps, flat_hand_mean,
                 mode='train'):

        self.model_type = model_type
        self.use_pca = use_pca
        self.num_pca_comps = num_pca_comps
        
        dataset_path = hydra.utils.to_absolute_path(dataset_path)
        self.betas = np.load(osp.join(dataset_path, 'mean_shape_{}.npy'.format(model_type)))
        self.gender = open(osp.join(dataset_path, 'gender.txt')).readlines()[0].strip()

        self.regstr_list = sorted(
            glob.glob(osp.join(dataset_path, mode, '*', model_type.upper(), '*.ply')))
        self.smplx_params_list = sorted(
            glob.glob(osp.join(dataset_path, mode, '*', model_type.upper(), '*.pkl')))
        self.pcl_list = sorted(
            glob.glob(osp.join(dataset_path, mode, '*', 'pcl', '*.ply')))
        
        assert len(self.regstr_list) == len(self.smplx_params_list) == len(self.pcl_list)

        if mode == 'test': # only use one sample for testing
            self.regstr_list = self.regstr_list[0:1]
            self.smplx_params_list = self.smplx_params_list[0:1]
            self.pcl_list = self.pcl_list[0:1]
        
        self.meta_info = {
            'gender': self.gender,
            'betas': self.betas,
            'representation': 'sdf',
            'use_pca': use_pca,
            'num_pca_comps': num_pca_comps,
            'flat_hand_mean': flat_hand_mean
        }

    def __getitem__(self, index):

        data = {}

        while True:
            try:
                regstr_mesh = trimesh.load(self.regstr_list[index])
                pcl = o3d.io.read_point_cloud(self.pcl_list[index])
                with open(self.smplx_params_list[index], 'rb') as f:
                    params_data = pkl.load(f, encoding='latin')
                    global_orient = params_data[
                        'global_orient'] if 'global_orient' in params_data else None
                    transl = params_data[
                        'transl'] if 'transl' in params_data else None
                    body_pose = params_data[
                        'body_pose'] if 'body_pose' in params_data else None
                    hand_pose = params_data[
                        'hand_pose'] if 'hand_pose' in params_data else None
                    left_hand_pose = params_data[
                        'left_hand_pose'] if 'left_hand_pose' in params_data else None
                    right_hand_pose = params_data[
                        'right_hand_pose'] if 'right_hand_pose' in params_data else None
                    leye_pose = np.zeros(3)
                    reye_pose = np.zeros(3)
                    jaw_pose = params_data[
                        'jaw_pose'] if 'jaw_pose' in params_data else None
                    betas = self.betas
                    expression = params_data[
                        'expression'] if 'expression' in params_data else None
                break
            except:
                print('corrupted ply/pkl')
                print(self.smplx_params_list[index])
                index = np.random.randint(self.__len__())

        # load SMPL(-H/X) registration
        regstr_verts = regstr_mesh.vertices - transl
        regstr_verts = Tensor(regstr_verts).float()
        regstr_faces = regstr_mesh.faces.astype(np.int32)
        regstr_faces = Tensor(regstr_faces).long()
        regstr_normals = Tensor(regstr_mesh.vertex_normals.copy()).float()
        data['regstr_verts'] = regstr_verts
        data['regstr_faces'] = regstr_faces
        data['regstr_normals'] = regstr_normals
        
        # load human point cloud
        pcl_verts = np.array(pcl.points) - transl
        pcl_verts = Tensor(pcl_verts).float()
        pcl_colors = Tensor(pcl.colors).float()
        pcl_normals = Tensor(pcl.normals).float()
        data['pcl_verts'] = pcl_verts
        data['pcl_colors'] = pcl_colors
        data['pcl_normals'] = pcl_normals
    
        # load SMPL(-H/X) parameters
        num_hand_pose = self.num_pca_comps if self.use_pca else 45

        if self.model_type == 'smplx':
            smpl_params = ms.ops.zeros([99 + 2 * num_hand_pose]).float()
            smpl_params[0] = 1
            smpl_params[4:7] = Tensor(global_orient).float()
            smpl_params[7:70] = Tensor(body_pose).float()
            smpl_params[70:70 +
                        num_hand_pose] = Tensor(left_hand_pose).float()
            smpl_params[70 + num_hand_pose:70 + 2 *
                        num_hand_pose] = Tensor(right_hand_pose).float()
            smpl_params[70 + 2 * num_hand_pose:73 +
                        2 * num_hand_pose] = Tensor(leye_pose).float()
            smpl_params[73 + 2 * num_hand_pose:76 +
                        2 * num_hand_pose] = Tensor(reye_pose).float()
            smpl_params[76 + 2 * num_hand_pose:79 +
                        2 * num_hand_pose] = Tensor(jaw_pose).float()
            smpl_params[79 + 2 * num_hand_pose:89 +
                        2 * num_hand_pose] = Tensor(betas).float()
            smpl_params[89 + 2 * num_hand_pose:99 +
                        2 * num_hand_pose] = Tensor(expression).float()
            data['smpl_thetas'] = smpl_params[7:70]
            data['smpl_exps'] = Tensor(expression).float()
        elif self.model_type == 'smplh':
            smpl_params = ms.ops.zeros([80 + 2 * num_hand_pose]).float()
            smpl_params[0] = 1
            smpl_params[4:7] = Tensor(global_orient).float()
            smpl_params[7:70] = Tensor(body_pose).float()
            smpl_params[70:70 +
                        num_hand_pose] = Tensor(left_hand_pose).float()
            smpl_params[70 + num_hand_pose:70 + 2 *
                        num_hand_pose] = Tensor(right_hand_pose).float()
            smpl_params[70 + 2 * num_hand_pose:80 +
                        2 * num_hand_pose] = Tensor(betas).float()
            data['smpl_thetas'] = smpl_params[7:70]
            data['smpl_exps'] = None
        elif self.model_type == 'smpl':
            smpl_params = ms.ops.zeros([86]).float()
            smpl_params[0] = 1
            smpl_params[4:7] = Tensor(global_orient).float()
            smpl_params[7:76] = Tensor(body_pose).float()
            smpl_params[76:86] = Tensor(betas).float()
            data['smpl_thetas'] = smpl_params[7:76]
            data['smpl_exps'] = None
        elif self.model_type == 'mano':
            smpl_params = ms.ops.zeros([17 + num_hand_pose]).float()
            smpl_params[0] = 1
            smpl_params[4:7] = Tensor(global_orient).float()
            smpl_params[7:7 + num_hand_pose] = Tensor(hand_pose).float()
            smpl_params[7 + num_hand_pose:17 +
                        num_hand_pose] = Tensor(betas).float()
            data['smpl_thetas'] = smpl_params[4:7 + num_hand_pose]
            data['smpl_exps'] = None
        else:
            raise NotImplementedError

        data['smpl_params'] = smpl_params
        return data

    def __len__(self):
        return len(self.smplx_params_list)